#! /usr/bin/python
import math
import ephi
from ephi import vect3d, Statics, ElectroDynamics, Scene, Screen;

statics = Statics()
dynamics = ElectroDynamics(statics)

# make a square coil cube polywell
wr = 0.025
side = 0.4
spacing = 0.01
current = 100000.0

hside = 0.5 * side
# due to spacing the real side is a bit expanded
hrside = 0.5 * (side + 2.0 * (2 * wr + spacing) / math.sqrt(2));

statics.addStaticElement (ephi.StaticRectangle(current, vect3d(-hrside, 0, 0),
    vect3d(0, -1, 0) * hside, vect3d(0, 0, -1) * hside, wr))
statics.addStaticElement (ephi.StaticRectangle(current, vect3d(hrside, 0, 0),
    vect3d(0, 1, 0) * hside, vect3d(0, 0, -1) * hside, wr))
statics.addStaticElement (ephi.StaticRectangle(current, vect3d(0, hrside, 0),
    vect3d(1, 0, 0) * hside, vect3d(0, 0, 1) * hside, wr))
statics.addStaticElement (ephi.StaticRectangle(current, vect3d(0, -hrside, 0),
    vect3d(1, 0, 0) * hside, vect3d(0, 0, -1) * hside, wr))
statics.addStaticElement (ephi.StaticRectangle(current, vect3d(0, 0, hrside),
    vect3d(-1, 0, 0) * hside, vect3d(0, 1, 0) * hside, wr))
statics.addStaticElement (ephi.StaticRectangle(current, vect3d(0, 0, -hrside),
    vect3d(1, 0, 0) * hside, vect3d(0, 1, 0) * hside, wr))

# render logarithmic bfield at two points
if 1:
  print 'Visualizing bfield..'
  scene = Scene(dynamics)
  screen = Screen(800, 800, 1, 1)
  scene.calc(screen)
  scene.set_coloring(0, True)
  scene.render_lic(screen)
  screen.write ('rectpoly_bfield1.bmp')
  screen.set_axis (screen.PLANE_XY, side / 3)
  scene.calc(screen)
  scene.set_coloring(32, True)
  scene.render_map(screen)
  screen.write ('rectpoly_bfield2.bmp')

# now we try to simulate a particle's path
if 1:
  print 'Simulating particle trajectory..'
  screen = Screen(800, 800, 0.8, 0.8)
  screen.draw_axes()
  statics.draw (screen)
  dynamics.setDTM(0.01)
  dynamics.inject_electron (vect3d(-side*0.7, side/32, 0), vect3d(1, 0, 0), 10000)
  for i in xrange(0,10000):
    screen.set_pixel (dynamics.get_pos(0), screen.RED)
    dynamics.step()
  screen.write ('rectpoly_sim.bmp')

# You can use this to render the scene in povray
# run: povray -geometry 800x800 +A +Q9 +AM2 polywell.pov
if 1:
  print 'Generating polywell.inc..'
  fp = open ('polywell.inc', 'w')
  fp.write (statics.get_povray_script())
  fp.close()

